import load as ld
import numpy as np
import matplotlib.pyplot as mp;

def cubicSplineInterpolationPrecal (X,Y):
    """compute the second derivate of Y in each point"""


    n= np.size(X)

    Ypp=np.zeros(n)
    # With equation 3.3.7(p115 Numerical recipes)
    # Ypp is the second derivatives 
    # A the matrix "tridiagonal" which represent coefficient of equation
    # B the vector
    # Thus we have the follow system 
    #       A * Ypp = b
    
    A = np.zeros([n,n])
    b = np.zeros(n)
    # Limit conditions authorize that y1"=0  and yn"=0   
    A[0,0] = 1 
    A[n-1,n-1] = 1
    b[0] = 0
    b[n-1] = 0
    
    for j in np.arange(1,n-1):
        A[j,j-1] = (X[j] - X[j-1])/6 
        A[j,j] = (X[j+1] - X[j-1])/3
        A[j,j+1] = (X[j+1] - X[j])/6
        b[j] = (Y[j+1] - Y[j])/(X[j+1] - X[j]) - (Y[j] - Y[j-1])/(X[j] - X[j-1])  
    # --- Solving the system ---

    Ypp = np.linalg.solve(A, b)
  
    return Ypp

 ##---------------------------------------------------------------------------------------------------------------##

def cubicSplineInterpolationCalc(X,Y,Ypp,x):
    
    N = X.size
    for i in np.arange(0,N-1):
        if (X[i] <= x) and (x < X[i+1]):
            # --- The 4 factors  ---
            A = (X[i+1] - x) / (X[i+1] - X[i])
            B = 1. - A
            tmp = 1./6.*(X[i+1]-X[i])**2
            C = tmp * (A**3 - A)
            D = tmp * (B**3 - B) 
           
            return A*Y[i] + B*Y[i+1] + C*Ypp[i] + D*Ypp[i+1]
        
    if x >= X[N-1]:
        return Y[N-1]
    else:
        return Y[0]

def cubicSplineInterpolation(X,Y):
    
    n = X.size
    # --- The second derivative in each point ---
    Ypp = cubicSplineInterpolationPrecal(X,Y)
    # --- Building the lambda function to return ---
    f = lambda x: cubicSplineInterpolationCalc(X,Y,Ypp,x)
    return f

##---------------------------------------------------------------------------------------------------------------##

def cubicSplineInterpolationDerivativeCalc(X,Y,Ypp,x):
    """ Returns the derivative of the cubic spline interpolation of the cloud of point X, Y, knowing its second derivative Ypp """
    
    n = X.size
    
    for i in np.arange(0,n-1):
        if (X[i] <= x) and (x < X[i+1]):
            # --- The 4 factors  ---
            A = (X[i+1] - x) / (X[i+1] - X[i])
            Ap = -1. / (X[i+1] - X[i])
            B = 1. - A
            Bp = 1. - Ap
            tmp = (1./6.) * (X[i+1] - X[i])**2
            C = (3.*Ap*(A**2) - Ap) * tmp
            D = (3.*Bp*(B**2) - Bp) * tmp
            
            return Ap*Y[i] + Bp*Y[i+1] + C*Ypp[i] + D*Ypp[i+1]
        
    if x >= X[n-1]:
        return Y[n-1]
    else:
        return Y[0]
    

def cubicSplineInterpolationDerivative(X,Y):
    """ Returns a f(x) function that is the derivative of the interpolation of the cloud of points X, Y using cubic Splines """

    N = X.size
    
    # --- The second derivative in each point ---
    Ypp = cubicSplineInterpolationPrecal(X,Y)
    
    # --- Building the lambda function to return ---
    return lambda x: cubicSplineInterpolationDerivativeCalc(X,Y,Ypp,x)


##---------------------------------------------------------------------------------------------------------------##


def cubicSplineInterpolationSecondDerivativeCalc(X,Y,Ypp,x):
    
    n = X.size
    
    for i in np.arange(0,n-1):
        if (X[i] <= x) and (x < X[i+1]):
            A = (X[i+1] - x) / (X[i+1] - X[i])
            B = 1. - A
            
            return A*Ypp[i] + B*Ypp[i+1]
        
    if x >= X[n-1]:
        return Y[n-1]
    else:
        return Y[0]
    
def cubicSplineInterpolationSecondDerivative(X,Y):

    N = X.size
    
    # --- The second derivative in each point ---
    Ypp = cubicSplineInterpolationPrecal(X,Y)
    
    # --- Building the lambda function to return ---
    return lambda x: cubicSplineInterpolationSecondDerivativeCalc(X,Y,Ypp,x)

